package io.github.eric.urpc.client;

import io.github.eric.urpc.core.conn.Connection;
import io.github.eric.urpc.core.conn.ConnectionListener;
import io.github.eric.urpc.core.conn.ConnectionManager;
import io.github.eric.urpc.core.exception.UrpcConnInitException;
import io.github.eric.urpc.core.exception.UrpcMsgSendException;
import io.github.eric.urpc.core.exception.UrpcRuntimeException;
import io.github.eric.urpc.core.rpc.Message;
import io.github.eric.urpc.core.rpc.Request;
import io.github.eric.urpc.core.rpc.Response;
import io.github.eric.urpc.core.rpc.msg.ApplicationInfoReq;
import io.github.eric.urpc.core.rpc.msg.ApplicationInfoRes;
import io.netty.util.Timeout;
import io.netty.util.TimerTask;
import org.apache.commons.pool2.impl.GenericObjectPool;
import org.apache.commons.pool2.impl.GenericObjectPoolConfig;

import java.net.InetSocketAddress;
import java.time.Duration;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.ReentrantLock;
import java.util.function.Predicate;
import java.util.stream.Collectors;

public class ClientConnectionManager implements ConnectionManager, ConnectionListener {
	private static final long RETRY_INTERVAL = 1000L;
	private static final int MAX_RETRIES = 1;
	private final ClientContext clientContext;
	private final ClientConnector clientConnector;
	private final Map<String, Connection> connectionMap;
	private final Map<String, PendingConnection> pendingConnections;
	private final GenericObjectPool<Connection> connectionPool;

	protected ClientConnectionManager(ClientContext clientContext, ClientConnector clientConnector) {
		this.clientContext = clientContext;
		this.clientConnector = clientConnector;
		watchConnection();
		connectionMap = new ConcurrentHashMap<>();
		pendingConnections = new ConcurrentHashMap<>();
		connectionPool = createConnectionPool();
	}

	@Override
	public void add(Connection conn) {
		Connection old = connectionMap.put(conn.getId(), conn);
		if (old == null) {
			clientContext.getLogger().info("[urpc client] on connection added: {}, local addr {}, remote addr {}",
					conn.getId(), conn.getLocalAddress(), conn.getRemoteAddress());
		}
	}

	@Override
	public void remove(Connection conn) {
		Connection old = connectionMap.remove(conn.getId());
		if (old != null) {
			clientContext.getLogger().info("[urpc client] on connection removed: {}, local addr {}, remote addr {}",
					old.getId(), old.getLocalAddress(), old.getRemoteAddress());
		}
	}

	@Override
	public void shutdown() {
		connectionPool.close();
		for (Map.Entry<String, Connection> entry : connectionMap.entrySet()) {
			Connection conn = entry.getValue();
			try {
				conn.close();
			} catch (Exception e) {
				clientContext.getLogger().error("[urpc client] close conn {} error, local addr {}, remote addr {}.",
						conn.getId(), conn.getLocalAddress(), conn.getRemoteAddress(), e);
			}
		}
	}

	public void validateConnectionInitStatus(Connection conn) {
		if (conn.isInitCompleted()) {
			return;
		}
		if (conn.isInitFailed()) {
			throw new UrpcConnInitException(String.format("[urpc client] conn %s init failed", conn.getId()));
		} else {
			PendingConnection pendingConnection = pendingConnections.get(conn.getId());
			if (pendingConnection == null) {
				throw new UrpcConnInitException(String.format("[urpc client] conn %s init in progress, but can't find pending lock", conn.getId()));
			}
			try {
				pendingConnection.getLock().lock();
				if (conn.isInitCompleted()) {
					return;
				}
				if (conn.isInitFailed()) {
					throw new UrpcConnInitException(String.format("[urpc client] conn %s init failed", conn.getId()));
				}
				long timeout = clientContext.getConnectionInitTimeout() * 1000L;
				long deadline = System.currentTimeMillis() + timeout;
				long remaining = timeout;
				while (remaining > 0) {
					try {
						pendingConnection.getCondition().await(remaining, TimeUnit.MILLISECONDS);
					} catch (InterruptedException e) {
						Thread.currentThread().interrupt();
					}
					if (!conn.isInitInProgress()) {
						break;
					}
					remaining = deadline - System.currentTimeMillis();
				}
				if (conn.isInitInProgress()) {
					throw new UrpcConnInitException(String.format("[urpc client] conn %s init timeout", conn.getId()));
				} else if (conn.isInitFailed()) {
					throw new UrpcConnInitException(String.format("[urpc client] conn %s init failed", conn.getId()));
				}
			} finally {
				pendingConnection.getLock().unlock();
			}
		}
	}

	/**
	 * can retry if async failed
	 * @param message
	 */
	public void sendAsync(Message message) {
		for (int attempt = 0; ; attempt++) {
			Connection conn = null;
			try {
				conn = connectionPool.borrowObject();
				conn.sendAsync(message);
				return;
			} catch (Exception e) {
				if (attempt < MAX_RETRIES) {
					if (conn != null) {
						clientContext.getLogger().warn("[urpc client] client conn manager async send msg {} failed, caused '{}', will retry after {}ms, conn {}, local addr {}, remote addr {}",
								message.getClass().getSimpleName(), e.getMessage(), RETRY_INTERVAL, conn.getId(), conn.getLocalAddress(), conn.getRemoteAddress());
					} else {
						clientContext.getLogger().warn("[urpc client] client conn manager async send msg {} failed, caused '{}', will retry after {}ms",
								message.getClass().getSimpleName(), e.getMessage(), RETRY_INTERVAL);
					}
				} else {
					if (conn != null) {
						throw new UrpcMsgSendException(
								String.format("[urpc client] client conn manager async send msg %s failed after retry, conn %s, local addr %s, remote addr %s",
										message.getClass().getSimpleName(), conn.getId(), conn.getLocalAddress(), conn.getRemoteAddress()), e);
					} else {
						throw new UrpcMsgSendException(String.format("[urpc client] client conn manager async send msg %s failed after retry",
								message.getClass().getSimpleName()), e);
					}
				}
				try {
					Thread.sleep(RETRY_INTERVAL);
				} catch (InterruptedException ie) {
					Thread.currentThread().interrupt();
					throw new UrpcMsgSendException(String.format("[urpc client] client conn manager async send msg %s interrupted",
							message.getClass().getSimpleName()));
				}
			} finally {
				if (conn != null) {
					connectionPool.returnObject(conn);
				}
			}
		}
	}

	/**
	 * can retry if sync failed
	 * @param request
	 */
	public <T extends Response> T sendSync(Request request) {
		for (int attempt = 0; ; attempt++) {
			Connection conn = null;
			try {
				conn = connectionPool.borrowObject();
				return conn.sendSync(request);
			} catch (Exception e) {
				if (attempt < MAX_RETRIES) {
					if (conn != null) {
						clientContext.getLogger().warn("[urpc client] client conn manager sync send request {} failed, caused '{}', will retry after {}ms, conn {}, local addr {}, remote addr {}",
								request.getClass().getSimpleName(), e.getMessage(), RETRY_INTERVAL, conn.getId(), conn.getLocalAddress(), conn.getRemoteAddress());
					} else {
						clientContext.getLogger().warn("[urpc client] client conn manager sync send request {} failed, caused '{}', will retry after {}ms",
								request.getClass().getSimpleName(), e.getMessage(), RETRY_INTERVAL);
					}
				} else {
					if (conn != null) {
						throw new UrpcMsgSendException(
								String.format("[urpc client] client conn manager sync send request %s failed after retry, conn %s, local addr %s, remote addr %s",
										request.getClass().getSimpleName(), conn.getId(), conn.getLocalAddress(), conn.getRemoteAddress()), e);
					} else {
						throw new UrpcMsgSendException(String.format("[urpc client] client conn manager sync send request %s failed after retry",
								request.getClass().getSimpleName()), e);
					}
				}
				try {
					Thread.sleep(RETRY_INTERVAL);
				} catch (InterruptedException ie) {
					Thread.currentThread().interrupt();
					throw new UrpcMsgSendException(String.format("[urpc client] client conn manager sync send request %s interrupted",
							request.getClass().getSimpleName()));
				}
			} finally {
				if (conn != null) {
					connectionPool.returnObject(conn);
				}
			}
		}
	}

	@Override
	public Collection<Connection> getConnections() {
		return Collections.unmodifiableCollection(connectionMap.values());
	}

	@Override
	public Collection<Connection> getConnections(Predicate<Connection> predicate) {
		if (predicate == null) {
			return getConnections();
		}
		Collection<Connection> connections = getConnections();
		return connections.stream().filter(predicate).collect(Collectors.toList());
	}

	@Override
	public void watchConnection() {
		if (clientContext.isShuttingDown()) {
			return;
		}
		try {
			clientContext.getHashedWheelTimer().newTimeout(new TimerTask() {
				public void run(Timeout timeout) throws Exception {
					try {
						int total = 0;
						Map<String, Integer> result = new HashMap<>();
						Collection<Connection> connections = getConnections();
						for (Connection conn : connections) {
							InetSocketAddress remoteAddress = conn.getRemoteAddress();
							if (remoteAddress == null) {
								int cur = result.computeIfAbsent("null", k -> 0);
								result.put("null", ++cur);
							} else {
								String host = remoteAddress.getHostString();
								int cur = result.computeIfAbsent(host, k -> 0);
								result.put(host, ++cur);
							}
							++total;
						}
						StringBuilder sb = new StringBuilder();
						sb.append("conn total: ").append(total);
						sb.append("\r\n");
						for (Map.Entry<String, Integer> entry : result.entrySet()) {
							sb.append("conn group by host '").append(entry.getKey()).append("': ").append(entry.getValue());
						}
						clientContext.getLogger().info("[urpc client] current client conn info\r\n{}", sb);
					} catch (Throwable t) {
						clientContext.getLogger().error("[urpc client] log client conn info error", t);
					} finally {
						watchConnection();
					}
				}
			}, 60, TimeUnit.SECONDS);
		} catch (Exception e) {
			clientContext.getLogger().error("[urpc client] create log client conn task error!", e);
		}
	}

	@Override
	public void onConnected(Connection conn) {
		add(conn);
		pendingConnections.put(conn.getId(), new PendingConnection(conn));
		// send application msg to server-side
		clientContext.getLogger().info("[urpc client] start to register application info to server side");
		ApplicationInfoReq req = new ApplicationInfoReq();
		req.setApplicationName(clientContext.getApplicationName());
		req.setApplicationId(clientContext.getApplicationId());
		conn.sendAsync(req);
	}

	@Override
	public void onAccepted(Connection conn) {
		throw new UnsupportedOperationException();
	}

	@Override
	public void onDisConnected(Connection conn) {
		remove(conn);
		reconnect();
	}

	public void reconnect() {
		try {
			clientContext.getHashedWheelTimer().newTimeout(new TimerTask() {
				public void run(Timeout timeout) throws Exception {
					// 如果是主动shutdown则不进行重连
					if (clientContext.isShuttingDown()) {
						return;
					}
					Connection conn = null;
					try {
						conn = connectionPool.borrowObject();
					} catch (Exception e) {
						reconnect();
					} finally {
						if (conn != null) {
							connectionPool.returnObject(conn);
						}
					}
				}
			}, 3, TimeUnit.SECONDS);
		} catch (Exception e) {
			InetSocketAddress remoteAddr = new InetSocketAddress(clientContext.getServerHost(), clientContext.getServerPort());
			clientContext.getLogger().error("[urpc client] create async connect task to addr {} failed!", remoteAddr, e);
		}
	}

	@Override
	public void updateConnectionInitStatus(String connId, boolean succeed) {
		Connection conn = connectionMap.get(connId);
		if (conn == null) {
			return;
		}
		PendingConnection pendingConnection = pendingConnections.get(connId);
		if (pendingConnection == null) {
			return;
		}
		try {
			pendingConnection.getLock().lock();
			if (succeed) {
				conn.setInitStatusCompleted();
			} else {
				conn.setInitStatusFailed();
			}
			pendingConnection.getCondition().signalAll();
		} catch (Exception e) {
			throw new UrpcRuntimeException(String.format("[urpc client] update conn %s init status error", conn.getId()), e);
		} finally {
			pendingConnection.getLock().unlock();
		}
	}

	private GenericObjectPool<Connection> createConnectionPool() {
		GenericObjectPoolConfig<Connection> config = new GenericObjectPoolConfig<>();
		config.setMaxTotal(clientContext.getConnectionPoolMaxTotal());
		config.setMaxIdle(clientContext.getConnectionPoolMaxIdle());
		config.setMinIdle(clientContext.getConnectionPoolMinIdle());
		config.setMaxWait(Duration.ofSeconds(clientContext.getConnectionPoolMaxWait()));
		config.setTestOnBorrow(true);
		config.setTimeBetweenEvictionRuns(Duration.ofSeconds(clientContext.getConnectionPoolTimeBetweenEvictionRuns()));
		config.setTestWhileIdle(true);

		ClientConnectionFactory factory = new ClientConnectionFactory(clientContext, clientConnector);
        return new GenericObjectPool<>(factory, config);
	}

	static class PendingConnection {
		private final Connection conn;
		private final ReentrantLock lock;
		private final Condition condition;

		public PendingConnection(Connection conn) {
			this.conn = conn;
			this.lock = new ReentrantLock();
			this.condition = lock.newCondition();
		}

		public Connection getConn() {
			return conn;
		}

		public ReentrantLock getLock() {
			return lock;
		}

		public Condition getCondition() {
			return condition;
		}
	}

	@Override
	public ApplicationInfoRes updateConnectionInitStatus(String connId, ApplicationInfoReq applicationInfoReq) {
		throw new UnsupportedOperationException();
	}
}
